import torch
from faster_whisper import WhisperModel
import os
import time
import opencc

def translate_audio(input_path, output_dir="translations", model_size="medium"):
    """
    使用 faster-whisper 翻译音频文件，并保存结果到文本文件

    参数:
        input_path (str): 输入的音频文件路径（如 .mp3）
        output_dir (str): 输出目录（默认 translations）
        model_size (str): 模型大小（默认 small，可选 tiny, base, small, medium, large-v1, large-v2）
    """
    print("CUDA 可用：", torch.cuda.is_available())
    print("当前使用 GPU：", torch.cuda.get_device_name(0))
    # os.environ["HF_HUB_OFFLINE"] = "1"
    converter = opencc.OpenCC('t2s')  # 't2s'表示繁体转简体
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    start_time = time.time()
    # 加载模型（自动下载或从缓存加载）
    model = WhisperModel(
        model_size,
        device="cuda",  # 自动选择 GPU（如果可用）或 CPU
        compute_type="float16",  # 量化类型（可选 float16, int8, int16）
        download_root = r"E:\backend-project\AI_Interview\python_AI_Interview_backend\model\faster-whisper-medium"  # 自定义模型缓存路径
    )

    print(f"实际加载模型：{model}")
    # 执行翻译
    segments, info = model.transcribe(
        input_path,
        beam_size=5,  # 解码束宽度（平衡速度与准确性）
        # language="zh",  # 可选：手动指定语言
        task="transcribe"
    )

    print(f"检测到语言: {info.language}，概率: {info.language_probability:.2f}")

    # 根据语言类型决定是否执行繁转简
    is_chinese = info.language == "zh"

    # 打印翻译结果（不写入文件）
    for segment in segments:
        text = segment.text
        if is_chinese:
            # 如果是中文进行额外处理
            text = converter.convert(text)  # 繁转简
        print(text)

    end_time = time.time()
    print("翻译完成！")
    print(f"翻译总耗时：{end_time - start_time:.2f}s")
if __name__ == "__main__":
    import argparse

    import torch

    # python download_fasterwhisper_model.py --input_path "E:\backend-project\AI_Interview\python_AI_Interview_backend\src\temp_audio_wav\1746717630564_20250513_125938.wav"

    parser = argparse.ArgumentParser(description="描述信息")
    parser.add_argument("input_path", help="音频文件路径")
    parser.add_argument("--output_dir", default="translations", help="输出目录")
    parser.add_argument("--model", default="medium", help="模型大小")
    args = parser.parse_args()


    translate_audio(
        input_path=args.input_path,
        output_dir=args.output_dir,
        model_size=args.model
    )
